import os

from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset
from torch.utils.data import SubsetRandomSampler
from torchvision.datasets import ImageFolder
from torchvision import transforms
from torchvision.transforms.transforms import GaussianBlur
from preprocess.transforms import ToTensor,Compose,RandomRotate,Center_Crop,RandomAffine

class TrainDataset(Dataset):
    def __init__(self, image_path_list):
        # generate image path list

        self.image_path_list = image_path_list

        self.transforms = transforms.Compose([
            # # transforms.Resize((224, 224)),
            # transforms.RandomApply( [transforms.RandomRotation(degrees=180)],p=0.5 ),
            # transforms.RandomAffine(degrees=10),
            # # transforms.CenterCrop((128, 128)),
            # transforms.RandomApply( [transforms.ColorJitter(brightness=0.5,contrast=0.5)],p = 0.15 ),
            # transforms.ToTensor(),

            # transforms.RandomHorizontalFlip(),
            # transforms.RandomVerticalFlip(),
            transforms.RandomRotation(degrees=180),
            transforms.RandomAffine(degrees=10, translate=(0.1,0.1),scale=(0.9,1.1)),
            # transforms.RandomAffine(degrees=10),
            transforms.CenterCrop(640),
            # transforms.ColorJitter(brightness=0.1,
            #                     #    contrast=0.1,
            #                     #    saturation=0.1
            #                        ),
            # transforms.RandomCrop((112,112)),
            transforms.ToTensor(),
            # transforms.RandomErasing(), # RandomErasing是对(c, h, w)形状的tensor进行操作，一般在ToTensor之后进行

        ])
 
        self.transforms_npy = Compose([
            ToTensor(),
            RandomRotate(30),
            RandomAffine(degrees=10, translate=(0.1,0.1)),
            Center_Crop(160),
            # ColorJitter(brightness=0.1,contrast=0.1,saturation=0.1)
        ])


    def __getitem__(self, index):
        image_path = self.image_path_list[index]
        
        if (image_path[-3:]=='npy'):
            image = np.load(image_path)
            label = int(image_path.split("/")[-2])
            label = torch.tensor(label)
            return self.transforms_npy(image).type(torch.FloatTensor), label
        else:
            image = Image.open(image_path)
            image = image .convert('RGB')
            # print (image_path,image.size,int(image_path.split("/")[-2]))
            # if image.size[0]!=640 or image.size[1]!=640 :
            #     print("-----------------")
            label = int(image_path.split("/")[-2])
            label = torch.tensor(label)
            return self.transforms(image), label

    def __len__(self):
        return len(self.image_path_list)
